# CF_MMX_functions with "macro" estimate of eta

#######################
# 1) "Meta-functions" #
#######################
# Ownership matrix
# n.b. the  & i!=j
build.co <- function(firm) {
  co <- matrix(0,nrow=length(firm),ncol=length(firm))
  for(i in 1:length(firm)) {
    for(j in 1:length(firm))
      if(firm[i]==firm[j] & i!=j) co[i,j] = 1  
  }
  return(co)
}
# New Grid maker: assigns models to firms, firms to countries.
# has option for random allocation using dirichlet-drawn sampling probabilities
## but DEFAULT allocation is a taking turns algorithm 
make.ownership <- function(n.firms,n.models){
  #
  f <- rep(1:n.firms,times=ceiling(n.models/n.firms))
  my.owner <- data.frame(m=1:n.models,f=f[1:n.models]) 
  #
  co.own.mat <- with(my.owner,build.co(f))
  #
  return(co.own.mat)
}
#
make.tradecosts <- function(n.countries,n.firms,n.models,v,settings){
#
    f <- rep(1:n.firms,times=ceiling(n.models/n.firms))
    my.owner <- data.frame(m=1:n.models,f=f[1:n.models]) 
#
    i <- rep(1:n.countries,times=ceiling(n.firms/n.countries))
    my.origin <- data.frame(f=1:n.firms,i=i[1:n.firms])  
#
  flow.frame <- expand.grid(m=1:n.models,n=1:n.countries)
  flow.frame <- dplyr::inner_join(flow.frame,my.owner,by="m")
  flow.frame <- dplyr::inner_join(flow.frame,my.origin,by="f")
  #trade costs
  bil.frame <- expand.grid(i=1:n.countries,n=1:n.countries)
  bil.frame <- within(bil.frame,{
     tau <- (1+ntb+tar) + abs(n-i)*settings$dave[v] # make.grid ain't v-specific, use dave1 in DGP.exog
     tau[i==n] <-1
  })
  flow.frame <- dplyr::inner_join(flow.frame,bil.frame,by=c("i","n"))
  return(data.frame(vtau=v,flow.frame))
}
##########################
#Exogenous part of DGP==== 
############ New version with shifters in income, Gini and quality ############
DGP.exog <- function(settings,v=1,grid){ # data generating process for the exogenous parameters and variable
  #settings are all the distributional parameters need to generate the data
  #grid are the mappings of countries, firms, and varieties
  ##Demand-side
  ##income parameters
  grid.hh  <- as.data.table(expand.grid(hh=1:n.hh,n=1:n.countries))
 # method of moments estimates of mu.ly and sigma.ly
 if(GINI==TRUE) {
  US.mn.y <- 48374/dol.units
  US.Gini <- 0.41
  grid.hh[, mn.y := US.mn.y * (1 + (n-1) * settings$income.slope[v]) ]
  grid.hh[, gini := US.Gini * (1 + (n-1) * settings$gini.slope[v]) ]
  grid.hh[, sigma.ly := sqrt(2)*qnorm((1+gini)/2) ]
  grid.hh[, mu.ly := log(mn.y) -0.5*sigma.ly^2 ]
  grid.hh[, y.h := exp(rnorm(n=n.hh*n.countries,mean=mu.ly,sd=sigma.ly)) ]
} else { #BLP estimates of mu.ly and sigma.ly
  grid.hh[, y.h := exp(rnorm(n=n.hh*n.countries,mean=MU.LY,sd=SIGMA.LY*sigma.ly.factor)) ] # BLP rep data gives the Trumpified moments
}
  grid.hh[, y.h.norm := log(median(y.h)) ] #bring to 0 to restore old results. this is over ALL consumers worldwide
  grid.hh[, muadd.a := settings$a2[v] * (log(y.h)-y.h.norm) ] 
  grid.hh[, alpha.h := exp(rnorm(n=n.hh*n.countries,mean=settings$mu.a[v]+muadd.a,sd=settings$sigma.a[v])) ] 
  grid.hh[, muadd.b := settings$b2[v] * (log(y.h)-y.h.norm)]  #setting 5 only; ISE via b.h
  b.h <- matrix(rnorm(n.hh*n.countries*n.x,mean=BETA[2:5] * settings$b.adj[v],sd= SIGMA.BETA[2:5] * settings$sigma.b.adj[v]),
                nrow=n.hh*n.countries,ncol=n.x,byrow=TRUE)
  
  
   # if there's an outside good
  if(U0==1) {
    b0.h <- rnorm(n.hh*n.countries,mean=BETA[1]+settings$mu.b0[v],sd=  SIGMA.BETA[1] * settings$sigma.b.adj[v])
    b.h <- cbind(b0.h,b.h)
    colnames(b.h) <- paste0("b",0:n.x,".h")
  } else {
    colnames(b.h) <- paste0("b",1:n.x,".h")
}
    demandside <- data.frame(grid.hh[,list(n,alpha.h,y.h)],b.h)  # n here is market
  #
  ##Supply-side
  grid.m <- grid[,.(f=first(f),i=first(i)),by=.(m)] # argument of DGP.exog() 
  grid.m[, mu.x := settings$mu.x[v] * (1 + (i-1) * settings$quality.slope[v]) ] 
  xandxi <- xandxi.BLP[vxi==v,.(x1,x2,x3,x4,xi,mc)] # keep relevant xi for this v [[add mc?]]
  xandxi <- as.matrix(xandxi[sample(.N,n.models)])  # where xandxi.BLP is a 131 X 5 data.table of the car attributes in BLP in 1990 and inferred xi and mc
  x <- xandxi[,1:4]  
  x0 <- rep(1,n.models) 
  # vector of cost determinants, cont's var logged
  w <- cbind(x0,log(x[,"x1"]),x[,"x2"],log(x[,"x3"]),log(x[,"x4"]))  # cost shifters
  colnames(w) <- paste0("w",0:n.x)
  # if there's an outside good
  if(U0==1) {
    x <- cbind(x0,x) # demand shifters
    colnames(x) <- paste0("x",0:n.x)
  } else {
    colnames(x) <- paste0("x",1:n.x)
  } 
  #random draws of xi that respect sd of the implied xi from BLP data
  xi <- xandxi[,5]
  mc <- xandxi[,6]
  #form data frame of the x and xi created earlier
  gw <- t(GAMMA %*% t(w))
  m.chars <- data.frame(m=1:n.models,x,gw,xi,mc)
  mf <- expand.grid(m=1:n.models,f=1:n.firms)
  mf <- dplyr::inner_join(mf,m.chars,by="m")
  # after sampling the mc, don't need omega or mc to be simulated since they were sampled 
  #merge back 
  supplyside <- dplyr::inner_join(grid,mf,by=c("m","f"))
  mc.range <<- range(supplyside$mc)
  #
  return(list(models=supplyside,households=demandside)) #comprises two data frames, 
      #one w/ n.models  rows and other w/ n.hh X n.countries rows
}


####################################
#Equilibrium p and s on each market==== 
####################################
eqbm <- function(dat,market,v=1,tar.change=0,verbose=FALSE) { 
# tar.change  = tar.post-tar.pre (less than 0 for tariff reductions)
# set verbose = TRUE if you want to see the Iter and Residual  for the FPITER.
# the new version of nov 2020 can deal with multiple X and outside good
# most important, it sources the main subs-functions. 
#
 source("eqbm_subfunctions.R",local=TRUE)
#  
  ## start of eqbm computation ##
  #models:
  df.m <- data.frame(subset(dat$models,n==market),tar.change=tar.change)
  df.m <- within(df.m,{tar.change[i==n]<-0})
  i <- df.m$i
  if(U0==0) xvars <- paste0("x",1:n.x) else xvars <- paste0("x",0:n.x)
  x <- as.matrix(df.m[,xvars]) # if df.m is only a data.frame
  colnames(x) <- xvars  # no effect for n.x>1, but names it x1 instead of just x
  xi <- df.m$xi
  mc <- df.m$mc * (df.m$tau + df.m$tar.change * df.m$tariff.change) # apply trade costs and name it mc, for consistency with subordinate functions
  #households:
  df <- subset(dat$households,n==market)
  if(U0==0) bvec <- paste0("b",1:n.x,".h") else bvec <- paste0("b",0:n.x,".h")
  b.h <- as.matrix(df[,bvec]) 
  colnames(b.h) <- bvec  # no effect for n.x>1, but names it b1.h instead of just b.h
  alpha.h <- df$alpha.h
  y.h <- df$y.h
  Y <- sum(df$y.h)
  co.own <- co.own.mat
  avg.opt.markup <- (mean(alpha.h)+1)/mean(alpha.h)
  eps <- -ownelas(prob.mat(mc*avg.opt.markup))
  p.monop.comp <- mc*(eps+1)/eps
  #eqbm:
  FPIresults <- fpiter(p.monop.comp,fixptfn=price.update,nu=settings.var$speed[v],control=list(trace=verbose,maxiter=MAXIT))
  n.loops  <- FPIresults$fpevals
  p <- FPIresults$par
  s <- share.mces(prob.mat(p))
  pelas.m <- ownelas(prm=prob.mat(p)) # model level own price elasticities
  avg.inc.m <- avg.inc(prob.mat(p))   #calculate average income per model
  avg.inc.results <- lm(log(avg.inc.m)~log(p))
  ISE <- as.numeric(avg.inc.results$coefficients[2])
  profit <- (p-mc)*((s*Y)/p) # firm-level profit (mc here is CIT, it has tau = 1 and tar.change = 0 for domestic flows)
  tarrev <- (tar+df.m$tar.change* df.m$tariff.change) * df.m$mc * ((s*Y)/p) #  (foreign) firm-level tariff rev on initial mc 
  #
  return(data.frame(m=df.m$m,f=df.m$f,i,n=market,x,xi,mc=df.m$mc,cit=mc,s,p,pelas.m,avg.inc.m,ISE,profit,tarrev,n.loops=n.loops,Y=Y))
}

#Putting all markets together, using eqbm() on each
DGP.endog <- function(parallel=TRUE,v=1,tar.change=0,data.exog){ #get equilibrium market shares and prices for each market and append them
if(parallel==T)  foreach(i=1:n.countries, .combine = 'rbind')  %dorng%  eqbm(dat=data.exog,i,v,tar.change) else foreach(i=1:n.countries, .combine = 'rbind')  %do%  eqbm(dat=data.exog,i,v,tar.change) 
}

#Mahalanobis distance caculations:
MahaDist <- function(data,features) {#data = a data frame  features= a character vector
MW <- as.data.table(data[,c("m",features)])    
ML <- melt(MW,id="m") # ML is the long version of MW (models, wide)
SD <- ML[,.(S_v=sd(value)),by=.(variable)]
X <- CJ(m=unique(MW$m),j=unique(MW$m),variable =SD$variable)  
X <- merge(X,ML,by=c("m","variable"))
setnames(X,c("value"),c("X_vm"))
X <- merge(X,ML,by.x = c("j","variable"),by.y=c("m","variable"))
setnames(X,c("value"),c("X_vj"))
X <- merge(X,SD,by="variable")
X[, sumterm := (X_vm-X_vj)^2/(S_v^2)]
X[S_v==0, sumterm := 0] # for any feature with S_v==0, we exclude from the sum
MD <- X[,.(maha_dist = sqrt(sum(sumterm))),by=.(m,j)]
  return(MD)
}
#Compensating variation==== 
#for MCES
CV <- function(dat,pre,post,market){
  # log.sum to calculate CS change
  log.sum <- function(p) {
    # matrix[n.hh X n.models] of indirect utilties: row= household, col = variety
    #CHANGES DONE
    ubar <- b.h %*% t(x) - alpha.h %*% t(log(p)) +rep(1,n.hh) %*% t(xi)
    # vector of household-level CS level
    log(U0+rowSums(exp(ubar)))
  }
  # price index : HH level  
  price.index <- function(logsum) {
    exp((-1/alpha.h)*logsum)
  }
  #
  # start of eqbm computation #  
  df.m <- data.frame(subset(dat$models,n==market))
  if(U0==0) xvars <- paste0("x",1:n.x) else xvars <- paste0("x",0:n.x)
  x <- as.matrix(df.m[,xvars]) # if df.m is only a data.frame
  colnames(x) <- xvars  # no effect for n.x>1, but names it x1 instead of just x
  xi <- df.m$xi
  #
  df <- subset(dat$households,n==market)
  if(U0==0) bvec <- paste0("b",1:n.x,".h") else bvec <- paste0("b",0:n.x,".h")
  b.h <- as.matrix(df[,bvec]) 
  colnames(b.h) <- bvec  # no effect for n.x>1, but names it b1.h instead of just b.h
  alpha.h <- df$alpha.h
  y.h <- df$y.h
  #
  logsum.pre <- log.sum(pre$p) #log sum term of expected indirect utility
  price.index.pre <- price.index(logsum.pre) #HH-level price index
  logsum.post <- log.sum(post$p) #log sum term of expected indirect utility
  price.index.post <- price.index(logsum.post) #HH-level price index
  #
  CV <- sum(y.h*(price.index.pre-price.index.post)/price.index.pre) # summing across individuals 
return(CV)
}

#for CES
CV.ces <- function(pre,post){
  CV.ces <- pre$Y*(1-post$phi.ces.eha^(-1/pre$eta)) # total change in CS
  return(CV.ces[1])
  }


#########################
#CES approximation   ====
#multi-market
estapprox.mm <- function(pre.input) { # using data on market shares and prices, costs and observable quality, 
  #estimate eta (and other parameters) under the CES approximation
  # CES market share #
  pre <-  as.data.table(pre.input)
  #
  share.ces.mm <- function(results) {
  term <-  pre[, term := exp(results$fitted.values)]$term
  sumterm_n <-  pre[, sumterm_n := sum(term), by=(n)]$sumterm_n
  share_in_n <-  pre[, share_in_n := sum(s), by=(n)]$share_in_n
  pre[,c("term","sumterm_n","share_in_n"):=NULL]
  return((term / sumterm_n)*share_in_n) 
  }
  #
  mc <- pre$cit
  tau <- pre$cit/pre$mc
  s <- pre$s
  p <- pre$p
  f <- pre$f
  m <- pre$m
  i <- pre$i
  n <- pre$n
  # estimate the demand and quality coefficient ("CES" and "B")
      ols.results <- lm(log(s)~log(tau)+factor(m)+factor(n)) 
      ptm.results <- lm(log(p)~log(tau)+factor(m)+factor(n)) 
      #
  cost.elas <- ols.results$coefficients[2]
  eta <- as.numeric(-cost.elas)
  ptm <- as.numeric(ptm.results$coefficients[2])
  p.ces <- mc*((eta+1)/eta)
  s.ces  <- share.ces.mm(results=ols.results)
  return(data.frame(pre,s.ces,p.ces,eta,ptm))
}
# aggregate version
estapproxagg.mm <- function(pre.input) { # using data on market shares and prices, costs and observable quality, 
  #estimate eta (and other parameters) under the CES approximation
  # CES market share #
  pre <-  as.data.table(pre.input)
  preagg <- pre[,.(s=sum(s),tau=mean(cit/mc)), by=.(i,n)]
  #
  tau <- preagg$tau
  s <- preagg$s
  i <- preagg$i
  n <- preagg$n
  # estimate the trade elas
  ols.results <- lm(log(s)~log(tau)+factor(i)+factor(n)) 
  #
  cost.elas <- ols.results$coefficients[2]
  eta <- as.numeric(-cost.elas)
  return(data.frame(pre,aggeta = eta))
}


##################################
# Counterfactual policy changes  ==== 

#exact hat algebra function for CES market share
share.EHA <- function(s.init,tau.hat,epsilon,S0) {
  (tau.hat^epsilon)/(S0+sum(s.init*tau.hat^epsilon))
}
#exact hat algebra of Phi
Phi.EHA <- function(s.init,tau.hat,epsilon,S0) {
  S0 + sum(s.init*tau.hat^epsilon)
}

# CES market share EHA
CF.approx <- function(pre,post,meth="EHA"){
  #  
  mc <- post$cit
  eta <- pre$eta[1]
  aggeta <- pre$aggeta[1]
  ptm <- pre$ptm[1]
  tau.hat <-  post$cit/pre$cit
  if(meth=="AHA") tau.hat <- 1+(tau.hat-1)*ptm
#
  outside.share <- 1 - sum(pre$s) # market share of the outside good, used in share.EHA() and Phi.EHA()
  outside.share.ces <- 1 - sum(pre$s.ces) # CES market share of the outside good, used in share.EHA() and Phi.EHA() for DEV
  if(meth=="DEV"){
    p.ces <- mc*((eta+1)/eta)
    s.ces  <- pre$s.ces * share.EHA(s.init=pre$s.ces,tau.hat=post$cit/pre$cit,epsilon=-eta,S0 = outside.share.ces)
    phi.ces.eha <- Phi.EHA(s.init=pre$s.ces,tau.hat=post$cit/pre$cit,epsilon=-eta,S0 = outside.share.ces)
  }
  else{
    p.ces <- pre$p * tau.hat
    s.ces  <- pre$s * share.EHA(s.init=pre$s,tau.hat=tau.hat,epsilon=-eta,S0 = outside.share)
    s.ces.agg  <- pre$s * share.EHA(s.init=pre$s,tau.hat=tau.hat,epsilon=-aggeta,S0 = outside.share)
    phi.ces.eha <- Phi.EHA(s.init=pre$s,tau.hat=tau.hat,epsilon=-eta,S0 = outside.share)
  }
  profit.ces <- (p.ces - mc)*((s.ces*pre$Y)/p.ces) # firm-level profit CES
  tar.change <- (tau.hat-1) * (pre$cit/pre$mc)
  tarrev.ces <- (tar+tar.change) * pre$mc *((s.ces*pre$Y)/p.ces)
  return(data.frame(post,s.ces,s.ces.agg,p.ces,phi.ces.eha,profit.ces,tarrev.ces))
}

#Compare pre and post situation
#CR5 function ()
CR5.old <- function(shares,firms) {
  x <- aggregate(cbind(firms,shares),by=list(firms),FUN=sum)$shares
  y <-sort(x,decreasing=T)
  return(sum(y[1:5]))  
}
CR5.agg <- function(DT) {
  setDT(DT[,c("s","f")])
  DT <- DT[,.(y = sum(s)),by=f][order(-y)]
  return(sum(DT$y[1:5]))  
}

#Comparison 
CF.compare <- function(pre,post,meth="EHA"){
  ## market share, considering insiders only:
  pre$s <- pre$s/sum(pre$s)  
  pre$s.ces <- pre$s.ces/sum(pre$s.ces)  
  post$s <- post$s/sum(post$s)  
  post$s.ces <- post$s.ces/sum(post$s.ces)  
  post$s.ces.agg <- post$s.ces.agg/sum(post$s.ces.agg)  
  ##
  #cr5 <- CR5.agg(shares=pre$s,firms=pre$f)
  cr5 <- CR5.agg(pre)
  cor.shr.pre <- cor(pre$s,pre$s.ces)
  reg.results <- lm(pre$s.ces~pre$s)
  reg.shr.lev <- as.numeric(reg.results$coefficients[2])
  #
  dom.shr <- with(pre,sum(s[i==n])) # initial domestic share
  s.delta <- post$s - pre$s   #true change
  change <- with(post,sum(s[i==n])) - dom.shr   #true domestic share change in pp
  #
  if(meth=="DEV"){
    dom.shr.ces <- with(pre,sum(s.ces[i==n])) # initial domestic share under CES
    change.ces <- with(post,sum(s.ces[i==n])) - dom.shr.ces 
    s.ces.delta <- post$s.ces - pre$s.ces
  }
  else{
    change.ces <- with(post,sum(s.ces[i==n])) - dom.shr 
    change.ces.agg <- with(post,sum(s.ces.agg[i==n])) - dom.shr 
    s.ces.delta <- post$s.ces - pre$s
  }
  #
  # compare sum of changes for domestic firms:
  cor.shr.chg <- cor(s.delta,s.ces.delta) # correlation across models in changes
  reg.results <- lm(s.ces.delta~s.delta) # RHS = "truth", slope >1 => CES magnifies i.e. overestimates changes for the largest true changers (domestic giants)
  reg.shr.dif <- as.numeric(reg.results$coefficients[2])
  #
  #PT
  pt.model <- (post$p[post$i!=post$n]-pre$p[pre$i!=pre$n])/(post$cit[post$i!=post$n]- pre$cit[pre$i!=pre$n])
  pte.model <- pt.model * (pre$cit[pre$i!=pre$n]/pre$p[pre$i!=pre$n])
  pt <- mean(pt.model[pt.model!= -Inf & pt.model!= Inf]) # pass thru rate (IO)
  pte <- mean(pte.model[pte.model!= -Inf & pte.model!= Inf]) #  pass thru elasticity  a la Int'l Macro (IM)
  s.f <- pre$s
  s.f[pre$i==pre$n] <- 0 #among foreign models
  s.f[post$cit-pre$cit==0] <-0 # among foreign models that have some variation in tariffs.
  r.f <- rank(-s.f)
  pte1 <- ((post$p[r.f==1]-pre$p[r.f==1])/(post$cit[r.f==1]- pre$cit[r.f==1]))*(pre$cit[r.f==1]/pre$p[r.f==1]) #  pass thru elasticity  a la Int'l Macro (IM) for top foreign firm NOT A MEAN
  #
  #Price elasticity (analytical)
  price.elas <- with(pre,mean(pelas.m))  # mean (across models) of own price elasticity
  #
  converge.prob <- max(pre$n.loops)==MAXIT | max(post$n.loops)==MAXIT
  #
  return(data.frame(eta=pre$eta[1],aggeta=pre$aggeta[1],ptm=pre$ptm[1],ISE=pre$ISE[1],cr5,pt,pte,pte1,price.elas,cor.shr.pre,reg.shr.lev,cor.shr.chg,reg.shr.dif,dom.shr,change,change.ces,change.ces.agg,converge.prob))
}

#############################
#Monte Carlo repetitions ==== 
#############################
#function needed for one rep 
monte.func.mm <- function(repi,settings,v=1,market=1,cf.meth="EHA") { # monte carlo stage
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
  #pre
  eqbm.pre.mm  <- as.data.table(DGP.endog(parallel=F,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm.with.ces <- as.data.table(estapprox.mm(eqbm.pre.mm)) #add ces approximation results
  eqbm.pre.mm.with.ces <- as.data.table(estapproxagg.mm(eqbm.pre.mm.with.ces)) #add aggregated ces approximation results
  #post:
  eqbm.post <- eqbm(dat=data.exog.mm,market=market,v=v,tar.change= 1) #compute eqbm on one market after removing tariffs
  eqbm.pre.with.ces <- eqbm.pre.mm.with.ces[n==market,] #keep only one destination
  eqbm.pre <- eqbm.pre.mm[n==market,] #keep only one destination
  #AFTER THIS, nothing in the code is changed compared to 2ctry, since we have only one destination left.
  eqbm.post.with.ces <- as.data.table(CF.approx(eqbm.pre.with.ces,eqbm.post,meth=cf.meth)) #add ces approximation results
  #
  return(data.frame(repi=repi,v=v,CF.compare(eqbm.pre.with.ces,eqbm.post.with.ces,meth=cf.meth)))
}

#Doing multiple monte carlo reps for a single market
monte.rep.mm <- function(parallel=TRUE,settings,v=1,market=1,cf.meth="EHA"){ #get equilibrium market shares and prices for each market and append them
  if(parallel==T)  foreach(i=1:n.reps, .combine = 'rbind')  %dorng%  monte.func.mm(i,settings,v,market,cf.meth) else foreach(i=1:n.reps, .combine = 'rbind')  %do%  monte.func.mm(i,settings,v,market,cf.meth) 
}

#Create output of monte for CF, this function is used everywhere: in MCES and MLOG, both MC and Oly panels.
monte.output.mm <- function(simdata,mix,pan){
simdata.out <- simdata[converge.prob==FALSE]
simdata.out <- simdata.out[,nobs:= 1:.N,by=v ] #counter 
simdata.out <- simdata.out[nobs <= n.reps.target] # keep only the first ones
#
saveRDS(simdata.out,file=paste("../Tables/CF_MM/",mix,"/monte_allv_",pan,"OG_",OGpct,"_aggeta.rds",sep="")) #Save this for graphs in particular
}  


